#
# This script plots data from M3L trace lines that contain:
# optimiz.c 595
#
# These lines are enabled in M3L by using the command:
# --print_bl_trace 1 --print_nni_scores 1
#
import math, os, sys, re

#
# CONSTANTS
#
#output_filename_seed = catchpath + ".cran"
title = "BL Optimization on the initial BIONJ tree"
xlab = "iterations"
ylab = "length change (subs/site)"
read_bfgs = False
bfgs_max_iter = 1000
first_topo = 0
max_topo = 10
YAXIS = "subtraction" # subtraction # division # raw



def normalize_lengths(data):
    for branch in data.keys():
        iters = data[branch].keys()
        iters.sort()
        
        starting_length = data[branch][ iters[0] ]
        if starting_length == 0.0:
            starting_length = 0.000000000000000001
        #print "Branch=", branch, "Starting length = ", starting_length
                
        for iter in iters:
            if YAXIS == "division":
                if (float(data[branch][iter]) == 0.0):
                    data[branch][iter] = 0.000000000000000001
                data[branch][iter] = math.log(float(data[branch][iter]) / starting_length)
            elif YAXIS == "subtraction":
                data[branch][iter] = float(data[branch][iter]) - starting_length
            elif YAXIS == "raw":
                data[branch][iter] = float(data[branch][iter])
    return data

def separate_unimax_iters(data):
    branches = data.keys()
    branches.sort()
    count_iter = 0
    for branch in branches:
        these_points = data[branch]
        these_iters = these_points.keys()
        these_iters.sort()
        data[branch] = {}
        for iter in these_iters:
            data[branch][count_iter] = these_points[ iter ]
            count_iter += 1
        #print "Branch", branch, data[branch]
    return data

def plot_lengths(data):
    strings = []
    maxx = None
    minx = 0.0
    maxy = None
    miny = None
    branches = data.keys()
    branches.sort()
    for branch in branches:
        iters = data[branch].keys()
        iters.sort()        
        string = "x" + branch.__str__() + " <-c("
        for x in iters:
            string += x.__str__() + ","   
            if maxx == None:
                maxx = x
            elif maxx < x:
                maxx = x
            if minx == None:
                minx = x
            elif minx > x:
                minx = x
            
        string = re.sub(",$", "", string)
        string += ");\n"
        
        string += "y" + branch.__str__() + "<-c("
        for x in iters:
            y = data[branch][x]
            string += y.__str__() + ","
            if maxy == None:
                maxy = y
            elif maxy < y:
                maxy = y
            if miny == None:
                miny = y
            elif miny > y:
                miny = y
                
        string = re.sub(",$", "", string)
        string += ");\n"
        string += "points(x" + branch.__str__() + ", y" + branch.__str__() + ", type='l', col='" + (branch+1).__str__() + "');\n"
        strings.append( string )
    
    string = "x <- c(" + minx.__str__() + "," + maxx.__str__() + ");\n"
    string += "y <- c(" + miny.__str__() + "," + maxy.__str__() + ");\n"
    log_string = ""
    #if YAXIS == "division":
    #    log_string = ", log='y'"
    string += "plot(x, y, type='n', main='" + title + "', xlab='" + xlab + "', ylab='" + ylab + "'" + log_string+ ");\n"
    for s in strings:
        string += s + "\n" 
    fout_cran = open(output_filename_seed + ".cran", "w")
    if read_bfgs == True:
        fout_cran.write("pdf('" + output_filename_seed + ".pdf', width=5, height=5);\n")
    else:
        fout_cran.write("pdf('" + output_filename_seed + ".pdf', width=15, height=5);\n")
    fout_cran.write(string)
    fout_cran.write("dev.off();\n")
    fout_cran.close()
    
    os.system("r --no-save < " + output_filename_seed + ".cran")


#
# get data.
#
def get_data(path, read_bfgs = False):
    
    fin = open(path, "r")
    lines = fin.readlines()
    fin.close()
        
    branch_iter_length = {}
    branch_swapscore = {}
    branch_lnl = {}
    lastiter = 0
    count_topos = 0
    found_nni_data = False
    for l in lines:
        if l.__contains__("[Topology"):
            count_topos += 1
            if count_topos > max_topo:
                break
        if l.__contains__(". NNI branch:") and count_topos == 1:
            found_nni_data = True
            tokens = l.split()
            branch = int(tokens[3])
            score = float(tokens[5])
            branch_swapscore[branch] = score
        if l.__contains__("optimiz.c 595"):
            if ((read_bfgs == False and count_topos == 1) or (read_bfgs == True and count_topos == 0)) and found_nni_data == False:
                print path, l
                tokens = l.split()
                branch = int(tokens[4])
                if branch < 0:
                    continue
                if (read_bfgs == True and l.__contains__("*")) or (read_bfgs == False):
                    length = float(tokens[8])
                    iter = int(tokens[6])
                    lnl = float(tokens[10])
                    branch_lnl[branch] = lnl
                    if iter > 0:
                        iter = 1
                    if branch not in branch_iter_length.keys():
                        branch_iter_length[branch] = {}
                    if iter != 0 or iter not in branch_iter_length[branch]:
                        branch_iter_length[branch][iter] = length
    return [branch_iter_length, branch_swapscore, branch_lnl]

#print branch_iter_length
#print branch_swapscore
#print branch_lnl


def print_branch_iter_length( umlen, mmlen, umswaps, mmswaps ):
    fout = open("bl_summary.txt", "w")
    fout.write("branch\tBIONJ_length\tUM_length\tMM_length\tUM_swap_score\tMM_swap_score\n")
    for i in mmlen.keys():
        if i in umlen:
            print i
            l = i.__str__() 
            l += "\t" + mmlen[i][0].__str__() #bionj
            l += "\t" + umlen[i][1].__str__()
            l += "\t" + mmlen[i][1].__str__()
            l += "\t" + umswaps[i].__str__()
            l += "\t" + mmswaps[i].__str__()
            print l
            fout.write( l + "\n")
    fout.close()
#
#            
#

umcatchpath = sys.argv[1]
mmcatchpath = sys.argv[2]

print "\n. getting data from UM. . ."
[umbranch_iter_length, umbranch_swapscore, umbranch_lnl] = get_data(umcatchpath)
print "\n. getting data from MM. . ."
[mmbranch_iter_length, mmbranch_swapscore, mmbranch_lnl] = get_data(mmcatchpath, read_bfgs = True)
print "\n. printing data. . ."
print_branch_iter_length(umbranch_iter_length, mmbranch_iter_length, umbranch_swapscore, mmbranch_swapscore)
#branch_iter_length = normalize_lengths( branch_iter_length )
#if (read_bfgs == False):
#    branch_iter_length = separate_unimax_iters( branch_iter_length )
#plot_lengths( branch_iter_length )

            